"""
把encoder和decoder合并形成seq2seq模型
"""

import torch.nn as nn

from chatbot.chatbot_core.decoder import Decoder
from chatbot.chatbot_core.encoder import Encoder


class Seq2seq(nn.Module):
    def __init__(self):
        super(Seq2seq, self).__init__()
        self.encoder = Encoder()
        self.decoder = Decoder()

    def forward(self, input, target, input_length, target_length):
        encoder_outputs, encoder_hidden = self.encoder(input, input_length)
        decoder_outputs, decoder_hidden = self.decoder(target, encoder_hidden, encoder_outputs)
        return decoder_outputs, decoder_hidden

    def evaluation(self, input, input_length):
        encoder_outputs, encoder_hidden = self.encoder(input, input_length)
        # decoded_sentence = self.decoder.evaluate(encoder_hidden, encoder_outputs)
        decoded_sentence = self.decoder.evaluate_beamsearch_heapq(encoder_hidden, encoder_outputs)
        return decoded_sentence
